import { type FC } from 'react';
import { useMemoizedFn, useSafeState, useMount, useUnmount } from 'ahooks';
import { SocketPro } from '@/libs/socket';
import Loading from '../Loading';
import SocketContext from './context';
import styles from './styles.module.less';

type SocketProps = {
  token: string;
  children?: any;
}

const wsUrl = import.meta.env.ENV_AI_SOCKET_URL;
const apiKey = import.meta.env.ENV_AI_SOCKET_APP_KEY;

const Socket: FC<SocketProps> = (props) => {
  const { children, token } = props;
  const [socket, setSocket] = useSafeState<SocketPro>();

  const createSocket = useMemoizedFn(() => {
    const s = new SocketPro(wsUrl, {
      apiKey,
      area: 'core',
      endpointId: 'openai',
      authorization: token
    });
    s.on('connect', () => setSocket(s));
  });

  useMount(createSocket);
  useUnmount(() => socket?.destroy());

  if (!socket) {
    return (
      <div className={styles.loading}>
        <Loading size={10}/>
      </div>
    );
  }

  return (
    <SocketContext.Provider value={{ socket }}>
      {children}
    </SocketContext.Provider>
  );
}

export default Socket;